# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : 泰坦尼克号生存案例.py
# @Author: dongguangwen
# @Date  : 2025-02-07 21:33
# 0.导入工具包
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import classification_report
from sklearn.tree import plot_tree
import matplotlib.pyplot as plt

# 1.加载数据
data = pd.read_csv('./data/train.csv')
# print(data.head())
# print(data.info())

# 2.数据处理
x = data[['Pclass', 'Sex', 'Age']].copy()
y = data['Survived'].copy()
# print(x.head(10))
x['Age'].fillna(x['Age'].mean(), inplace=True)
# print(x.head(10))
x = pd.get_dummies(x)  # 需要转数值one-hot编码

x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=22)

# 3.模型训练
model = DecisionTreeClassifier()
model.fit(x_train, y_train)

# 4.模型评估
y_pred = model.predict(x_test)
# print(y_pred)
print(classification_report(y_test, y_pred))

# 5.可视化
# plot_tree(model)
# plt.show()

plt.figure(figsize=(30, 20))
plot_tree(model, max_depth=10, filled=True, feature_names=['Pclass', 'Age', 'Sex_female', 'Sex_male'], class_names=['died', 'survived'])
plt.show()

"""
              precision    recall  f1-score   support

           0       0.80      0.86      0.83       110
           1       0.75      0.65      0.70        69

    accuracy                           0.78       179
   macro avg       0.77      0.76      0.76       179
weighted avg       0.78      0.78      0.78       179
"""
